%This scrip organizes the estimation moments from FGKK for accounting for parameter uncertainty in the test.
%Corresponding author: Rodrigo Adao
%Date: 09/11/2024
%Input: estimation_data_fgkk.mat
%Outut: output_estimation_step.mat

%% Preliminaries

clear;
close all;

%set lacal paths
function_path = "..\Functions\";
addpath(function_path,'-frozen');
set_local_paths;

%% Import and adjust data from estimation stage

%Estimation of sigma, omega_star and sigma_star
    load(data_path + "estimation_data_fgkk.mat") 

%Number of obs for estimation of each parameter in FGKK
    Nsigma  = 2454023;  %Tab 4
    Nomegas = 2454023;  %Tab 4
    Nsigmas = 2564731;  %Tab 7    
    Neta    = 371916;   %Tab 6
    Nkappa  = 2041;     %Tab 5

%Select varietes 
    %Since we only consider summations over objects multiplied by tariff
    %changes, it's wlog to select the subset of varieties with non-zero tariff changes

    %Import varieties
    samp_dtau = sum(abs(omega_sigma_dtf), 2) > 0 ;     %sample of import varieties used in estimation with non-zero tariff changes
    
    diff_tau = omega_sigma_diff_tf(samp_dtau ,:);
    delta_tau = omega_sigma_dtf(samp_dtau ,:);
    
    sigma_dlpduty_resid = sigma_dlpduty_resid(samp_dtau,:);
    omega_star_dlq1_resid = omega_star_dlq1_resid(samp_dtau,:);
    
    sigma_R  = sigma_resid(samp_dtau,:)/Nsigma; 
    omegas_R = omega_star_resid(samp_dtau,:)/Nomegas;    
    eta_R    = eta_R(samp_dtau,:);
    kappa_R  = kappa_R(samp_dtau,:);
    
    all_m_variety_tab = [all_vty(samp_dtau,:), all_vty_rownum(samp_dtau,:), all_vty_colnum(samp_dtau,:)];
    
    %Export varieties
    samp_dtau_star = sum(abs(sigma_star_dtf_star), 2) > 0 ;    %sample of export varieties used in estimation with non-zero tariff changes
    
    diff_tau_star = sigma_star_diff_tf_star(samp_dtau_star ,:);
    delta_tau_star = sigma_star_dtf_star(samp_dtau_star,:);
    sigma_star_dlpduty_resid = sigma_star_dlpduty_resid(samp_dtau_star,:);
    
    sigmas_R = sigma_star_resid(samp_dtau_star,:)/Nsigmas; 

    all_x_variety_tab = [all_vty(samp_dtau_star,:), all_vty_rownum(samp_dtau_star,:), all_vty_colnum(samp_dtau_star,:)];

%% Build Lambda vectors for each parameter
    Lambda_sigma  = delta_tau.*sigma_R;
    Lambda_omegas = delta_tau.*omegas_R;
    Lambda_sigmas = delta_tau_star.*sigmas_R;   
    Lambda_eta    = delta_tau.*eta_R;
    Lambda_kappa  = delta_tau.*kappa_R;

%% Create jacobian of estimation moments
    %Jacobian of sigma moment
    Jgamma_sigma_sigma = sum( delta_tau.*sigma_dlpduty_resid , 'all')/Nsigma;
    
    %Jacobian of omega_star moment
    Jgamma_omegas_omegas = - sum( delta_tau.*omega_star_dlq1_resid , 'all')/Nomegas;
    
    %Jacobian of sigma_star moment
    Jgamma_sigmas_sigmas = sum( delta_tau_star.*sigma_star_dlpduty_resid , 'all')/Nsigmas;
    
    %Jacobian of eta moment
    Jgamma_eta_eta   = est_eta_resids.eta_dlstattf1'*est_eta_resids.eta_dlpindex_resid/Neta;
    Jgamma_eta_sigma = est_eta_resids.eta_dlstattf1'*est_eta_resids.eta_v_resid/Neta; 
    
    %Jacobian of kappa moment
    Jgamma_kappa_kappa = est_kappa_resids.kappa_dlstattf1'*est_kappa_resids.kappa_dlpindex_resid/Nkappa;
    Jgamma_kappa_sigma = est_kappa_resids.kappa_dlstattf1'*est_kappa_resids.kappa_dlpindex_dsigma_resid/Nkappa; 
    Jgamma_kappa_eta   = est_kappa_resids.kappa_dlstattf1'*est_kappa_resids.kappa_dlpindex_deta_resid/Nkappa; 
    
    %Build 5 x 5 jacobian matrix of estimation moments
    Jgamma = [Jgamma_sigma_sigma, 0                     , 0                   , 0               , 0                 ;
              0                 , Jgamma_omegas_omegas  , 0                   , 0               , 0                 ;
              0                 , 0                     , Jgamma_sigmas_sigmas, 0               , 0                 ;
              Jgamma_eta_sigma  , 0                     , 0                   , Jgamma_eta_eta  , 0                 ;
              Jgamma_kappa_sigma, 0                     , 0                   , Jgamma_kappa_eta, Jgamma_kappa_kappa];


%% Compute the variance matrix estimation moments

    %Terms associated with import varieties
    NMs = length(delta_tau);
    
    %diagonal terms
    Vgamma_sigma_sigma   = zeros(NMs,1);
    Vgamma_omegas_omegas = zeros(NMs,1);
    Vgamma_eta_eta       = zeros(NMs,1);
    Vgamma_kappa_kappa   = zeros(NMs,1);
    
    for c = 1:NMs
        Vgamma_sigma_sigma(c)    = sum( Lambda_sigma(c,:)' *Lambda_sigma(c,:) , 'all');
        Vgamma_omegas_omegas(c)  = sum( Lambda_omegas(c,:)'*Lambda_omegas(c,:), 'all');    
        Vgamma_eta_eta(c)        = sum( Lambda_eta(c,:)'   *Lambda_eta(c,:)   , 'all');    
        Vgamma_kappa_kappa(c)    = sum( Lambda_kappa(c,:)' *Lambda_kappa(c,:) , 'all');    
    end

    %Off-diagonal terms
    Vgamma_sigma_omegas  = zeros(NMs,1);
    Vgamma_sigma_eta   = zeros(NMs,1);
    Vgamma_sigma_kappa   = zeros(NMs,1);
    
    Vgamma_omegas_eta   = zeros(NMs,1);
    Vgamma_omegas_kappa   = zeros(NMs,1);
    
    Vgamma_eta_kappa   = zeros(NMs,1);

    for c = 1:NMs       
        Vgamma_sigma_omegas(c)   = sum( Lambda_omegas(c,:)'*Lambda_sigma(c,:) , 'all');    
        Vgamma_sigma_eta(c)      = sum( Lambda_eta(c,:)'   *Lambda_sigma(c,:) , 'all');   
        Vgamma_sigma_kappa(c)    = sum( Lambda_kappa(c,:)' *Lambda_sigma(c,:) , 'all');   
    
        Vgamma_omegas_eta(c)     = sum( Lambda_eta(c,:)'   *Lambda_omegas(c,:) , 'all');   
        Vgamma_omegas_kappa(c)   = sum( Lambda_kappa(c,:)' *Lambda_omegas(c,:) , 'all');   
        
        Vgamma_eta_kappa(c)      = sum( Lambda_kappa(c,:)' *Lambda_eta(c,:)    , 'all');   
    end

    %Terms associated with export varieties
    NXs = length(delta_tau_star);
    Vgamma_sigmas_sigmas = zeros(NXs,1);
    for c = 1:NXs
        Vgamma_sigmas_sigmas(c)  = sum( Lambda_sigmas(c,:)'*Lambda_sigmas(c,:) , 'all');  
    end

    %Build 5 x 5 matrix
    Vgamma = [sum(Vgamma_sigma_sigma) , sum(Vgamma_sigma_omegas) ,   0                      , sum(Vgamma_sigma_eta), sum(Vgamma_sigma_kappa)  ;
              sum(Vgamma_sigma_omegas), sum(Vgamma_omegas_omegas),   0                      , sum(Vgamma_omegas_eta), sum(Vgamma_omegas_kappa);
              0                       , 0                        , sum(Vgamma_sigmas_sigmas), 0                     , 0                       ;
              sum(Vgamma_sigma_eta)   , sum(Vgamma_omegas_eta)   ,   0                      , sum(Vgamma_eta_eta)   , sum(Vgamma_eta_kappa)   ;
              sum(Vgamma_sigma_kappa) , sum(Vgamma_omegas_kappa) ,   0                      , sum(Vgamma_eta_kappa) , sum(Vgamma_kappa_kappa) ];


%% Select subset of varieties that are considered for counterfactuals

    testing_Mshifts = (all_m_variety_tab(:,3) == -99) + (all_m_variety_tab(:,4) == -99) == 0;
    all_m_variety_tab = all_m_variety_tab(testing_Mshifts , :);
    diff_tau          = diff_tau(testing_Mshifts , :);
    delta_tau         = delta_tau(testing_Mshifts , :);
    Lambda_sigma      = Lambda_sigma(testing_Mshifts , :);
    Lambda_omegas     = Lambda_omegas(testing_Mshifts , :);
    Lambda_eta        = Lambda_eta(testing_Mshifts , :);
    Lambda_kappa      = Lambda_kappa(testing_Mshifts , :);
    
    testing_Xshifts = (all_x_variety_tab(:,3) == -99) + (all_x_variety_tab(:,4) == -99) == 0;
    all_x_variety_tab = all_x_variety_tab(testing_Xshifts , :);
    diff_tau_star  = diff_tau_star(testing_Xshifts , :);
    delta_tau_star = delta_tau_star(testing_Xshifts , :);
    Lambda_sigmas  = Lambda_sigmas(testing_Xshifts , :);
    
    diff_tau_est = diff_tau;
    diff_tau_star_est = diff_tau_star;
    delta_tau_est = delta_tau;
    delta_tau_star_est = delta_tau_star;

%% Save output      
save(save_estimation_path + "output_estimation_step.mat", 'Vgamma', 'Jgamma', 'Lambda_sigma', 'Lambda_omegas', 'Lambda_sigmas', 'Lambda_eta', 'Lambda_kappa', 'delta_tau_est', 'delta_tau_star_est', 'diff_tau_est', 'diff_tau_star_est', 'all_x_variety_tab', 'all_m_variety_tab');
